package com.ideaaedi.commonds.security;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA公钥/私钥 工具类
 *
 * @author <font size = "20" color = "#3CAA3C"><a href="https://gitee.com/JustryDeng">JustryDeng</a></font> <img src="https://gitee.com/JustryDeng/shared-files/raw/master/JustryDeng/avatar.jpg" />
 * @since  2020/5/10 14:48:13
 */
@Slf4j
@SuppressWarnings("unused")
public final class RSA extends BaseSecuritySupport {
    
    private RSA() {
        throw new UnsupportedOperationException("util class cannot supported instance");
    }

    private static final String KEY_ALGORITHM = "RSA";
    private static final String SIGN_ALGORITHM = "SHA256withRSA";
    private static final int KEY_SIZE = 2048;
    private static final String PUBLIC_KEY = "RSAPublicKey";
    private static final String PRIVATE_KEY = "RSAPrivateKey";
    
    /**
     * 公钥加密
     *
     * @param content
     *         需要被加密的内容
     * @param publicKeyStr
     *         公钥
     *
     * @return 加密后的内容
     */
    public static String encrypt(String content, String publicKeyStr) throws NoSuchAlgorithmException,
            InvalidKeySpecException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException,
            InvalidKeyException, IOException {
        RSAPublicKey publicKey = (RSAPublicKey)loadPublicKeyByStr(publicKeyStr);
        Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        ///if (log.isDebugEnabled()) {
        ///    log.debug("cipher algorithm -> {}", cipher.getAlgorithm());
        ///}
        byte[] contentBytes = content.getBytes(StandardCharsets.UTF_8);
        int contentLength = contentBytes.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        // 分段加密加密块最大大小 （注：因为PKCS1Padding占用11个字节，所以这里减去11）
        int maxEncryptBlockSize = publicKey.getModulus().bitLength() / 8 - 11;
        int offset = 0;
        byte[] cache;
        while (offset < contentLength) {
            int inputLen = Math.min(maxEncryptBlockSize, contentLength - offset);
            cache = cipher.doFinal(content.getBytes(), offset, inputLen);
            out.write(cache, 0, cache.length);
            offset += inputLen;
        }
        byte[] encryptedData = out.toByteArray();
        out.close();
        return Base64.getEncoder().encodeToString(encryptedData);
    }
    
    /**
     * 公钥加密
     *
     * @param content
     *         需要被加密的内容
     * @param publicKeyStr
     *         公钥
     *
     * @return 加密后的内容
     */
    @SneakyThrows
    public static String encryptSilently(String content, String publicKeyStr) {
        return RSA.encrypt(content, publicKeyStr);
    }
    
    /**
     * 私钥解密
     *
     * @param content
     *         需要被解密的内容
     * @param privateKeyStr
     *         私钥
     *
     * @return 解密后的内容
     */
    public static String decrypt(String content, String privateKeyStr) throws NoSuchAlgorithmException, InvalidKeySpecException,
            NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, IOException {
        RSAPrivateKey privateKey = (RSAPrivateKey)loadPrivateKeyByStr(privateKeyStr);
        byte[] contentBytes = Base64.getDecoder().decode(content.getBytes(StandardCharsets.UTF_8));
        Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
        ///if (log.isDebugEnabled()) {
        ///    log.debug("cipher algorithm -> {}", cipher.getAlgorithm());
        ///}
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        int contentLength = contentBytes.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        // RSA密文块大小（ = 密钥长度/8 ）
        int maxDecryptBlockSize = privateKey.getModulus().bitLength() / 8;
        int offset = 0;
        byte[] cache;
        while (offset < contentLength) {
            cache = cipher.doFinal(contentBytes, offset, maxDecryptBlockSize);
            out.write(cache, 0, cache.length);
            offset += maxDecryptBlockSize;
        }
        out.close();
        return out.toString(StandardCharsets.UTF_8);
    }
    
    /**
     * 私钥解密
     *
     * @param content
     *         需要被解密的内容
     * @param privateKeyStr
     *         私钥
     *
     * @return 解密后的内容
     */
    @SneakyThrows
    public static String decryptSilently(String content, String privateKeyStr) {
        return RSA.decrypt(content, privateKeyStr);
    }

    /**
     * 分段加密/解密通用方法
     */
    private static byte[] doFinal(Cipher cipher, byte[] data, int blockSize) throws IllegalBlockSizeException, BadPaddingException, IOException {
        int inputLen = data.length;
        int offSet = 0;
        byte[][] blocks = new byte[(inputLen / blockSize) + 1][];
        for (int i = 0; offSet < inputLen; i++) {
            int len = Math.min(blockSize, inputLen - offSet);
            blocks[i] = cipher.doFinal(data, offSet, len);
            offSet += blockSize;
        }
        // 合并结果
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] block : blocks) {
            out.write(block);
        }
        return out.toByteArray();
    }

    /**
     * 签名
     */
    public static String sign(String privateKeyStr, String content)
            throws NoSuchAlgorithmException, InvalidKeyException, SignatureException, InvalidKeySpecException {
        RSAPrivateKey privateKey = (RSAPrivateKey)loadPrivateKeyByStr(privateKeyStr);
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initSign(privateKey);
        signature.update(content.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(signature.sign());
    }

    /**
     * 签名
     */
    @SneakyThrows
    public static String signSilently(String privateKeyStr, String content){
        return sign(privateKeyStr, content);
    }

    /**
     * 验签
     */
    public static boolean verify(String publicKeyStr, String content, String sign)
            throws NoSuchAlgorithmException, InvalidKeyException, SignatureException, InvalidKeySpecException {
        PublicKey publicKey = loadPublicKeyByStr(publicKeyStr);
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initVerify(publicKey);
        signature.update(content.getBytes(StandardCharsets.UTF_8));
        return signature.verify(Base64.getDecoder().decode(sign));
    }

    /**
     * 验签
     */
    @SneakyThrows
    public static boolean verifySilently(String publicKeyStr, String content, String sign){
        return verify(publicKeyStr, content, sign);
    }

    /**
     * 生成RSA公私钥
     *
     * @param format 指定公私钥格式
     *               <br />
     *               注：虽然参数是数组，但是只识别数组中的第一个元素
     *               <br />
     *               目前支持的格式有：
     *               <ul>
     *                   <li>默认格式（无参）：无----BEGIN xxx KEY-----开头，无----END xxx KEY-----结尾，无换行的格式</li>
     *                   <li>PEM (不区分大小写)：有----BEGIN xxx KEY-----开头，有----END xxx KEY-----结尾，有换行的格式</li>
     *               </ul>
     *
     * @return <ul>
     * <li>左-公钥</li>
     * <li>右-私钥</li>
     * </ul>
     */
    public static Pair<String, String> generateRsaKeyPair(String... format) throws NoSuchAlgorithmException {
        Map<String, Object> keyMap = initKey();
        Key publicKey = (Key) keyMap.get(PUBLIC_KEY);
        String publicKeyStr = Base64.getEncoder().encodeToString(publicKey.getEncoded());
        Key privateKey = (Key) keyMap.get(PRIVATE_KEY);
        String privateKeyStr = Base64.getEncoder().encodeToString(privateKey.getEncoded());

        // 转换为PEM格式
        if (format != null && format.length > 0 && "PEM".equalsIgnoreCase(format[0])) {
            publicKeyStr = "-----BEGIN PUBLIC KEY-----\n" +
                    publicKeyStr.replaceAll("(.{64})", "$1\n") +
                    "\n-----END PUBLIC KEY-----";
            privateKeyStr = "-----BEGIN PRIVATE KEY-----\n" +
                    privateKeyStr.replaceAll("(.{64})", "$1\n") +
                    "\n-----END PRIVATE KEY-----";
        }
        return Pair.of(publicKeyStr, privateKeyStr);
    }

    /**
     * 生成RSA公私钥
     *
     * @param format 指定公私钥格式
     *               <br />
     *               注：虽然参数是数组，但是只识别数组中的第一个元素
     *               <br />
     *               目前支持的格式有：
     *               <ul>
     *                   <li>默认格式（无参）：无----BEGIN xxx KEY-----开头，无----END xxx KEY-----结尾，无换行的格式</li>
     *                   <li>PEM (不区分大小写)：有----BEGIN xxx KEY-----开头，有----END xxx KEY-----结尾，有换行的格式</li>
     *               </ul>
     *
     * @return <ul>
     * <li>左-公钥</li>
     * <li>右-私钥</li>
     * </ul>
     */
    @SneakyThrows
    public static Pair<String, String> generateRsaKeyPairSilently(String... format) {
        return RSA.generateRsaKeyPair(format);
    }
    
    /**
     * 将字符串私钥转换为RSAPrivateKey
     *
     * @param privateKeyStr
     *         字符串私钥
     *
     * @return RSAPrivateKey实例
     */
    private static PrivateKey loadPrivateKeyByStr(String privateKeyStr) throws NoSuchAlgorithmException,
            InvalidKeySpecException {
        /*
         * 1. 移除字符串中的"-----BEGIN PUBLIC KEY-----"部分
         * 2. 移除字符串中的"-----BEGIN PRIVATE KEY-----"部分
         * 3. 使用正则表达式\\s+匹配一个或多个空白字符（包括空格、制表符、换行符等），并将其替换为空字符串""，即删除所有空白字符
         */
        privateKeyStr = privateKeyStr.replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s+", "");
        byte[] encodedKey = Base64.getDecoder().decode(privateKeyStr);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encodedKey);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }
    
    /**
     * 将字符串公钥转换为RSAPublicKey
     *
     * @param publicKeyStr
     *         字符串公钥
     * @throws NoSuchAlgorithmException 当算法不支持时抛出
     * @throws InvalidKeySpecException 当key无效时抛出
     * @return RSAPublicKey实例
     */
    private static PublicKey loadPublicKeyByStr(String publicKeyStr)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        /*
         * 1. 移除字符串中的"-----BEGIN PUBLIC KEY-----"部分
         * 2. 移除字符串中的"-----END PUBLIC KEY-----"部分
         * 3. 使用正则表达式\\s+匹配一个或多个空白字符（包括空格、制表符、换行符等），并将其替换为空字符串""，即删除所有空白字符
         */
        publicKeyStr = publicKeyStr.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s+", "");
        byte[] encodedPublicKey = java.util.Base64.getDecoder().decode(publicKeyStr);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedPublicKey);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }
    
    /**
     * 初始化key
     */
    private static Map<String, Object> initKey() throws NoSuchAlgorithmException {
        // 获取密钥对生成器
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        // 设置密钥位数
        keyPairGen.initialize(KEY_SIZE, new SecureRandom());
        // 生成密钥对
        KeyPair keyPair = keyPairGen.generateKeyPair();
        // 公钥
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        // 私钥
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        Map<String, Object> keyMap = new HashMap<>(4);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }
    
}